package com.ds.infrastructure.hyperspace.container.internal.opentracing.impl;

import com.ds.infrastructure.hyperspace.container.internal.util.Http2Util;
import io.netty.handler.codec.http.HttpRequest;
import io.netty.handler.codec.http2.Http2Headers;
import org.apache.commons.lang3.StringUtils;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import static com.ds.infrastructure.hyperspace.container.internal.opentracing.impl.TracerConstants.*;

/**
 * @Author: yaozh
 * @Description:
 */
public class HyperspaceTracerUtil {
    private HyperspaceTracerUtil() {
    }

    public static TracerSpan extract(HttpRequest httpRequest) {
        String traceId = null;
        String spanId = null;
        Map<String, String> baggage = new HashMap<>();
        Iterator<Map.Entry<String, String>> entries = httpRequest.headers().entries().iterator();
        while (entries.hasNext()) {
            Map.Entry<String, String> entry = entries.next();
            if (StringUtils.equalsIgnoreCase(BAGGAGE_ZIPKIN_TRACE_ID_KEY, entry.getKey())) {
                traceId = entry.getValue();
            } else if (StringUtils.equalsIgnoreCase(BAGGAGE_ZIPKIN_SPAN_ID_KEY, entry.getKey())) {
                spanId = entry.getValue();
            } else if (StringUtils.equalsIgnoreCase(BAGGAGE_KEY_ZIPKIN_SAMPLED, entry.getKey())) {
                baggage.put(BAGGAGE_KEY_ZIPKIN_SAMPLED, entry.getValue());
            } else if (StringUtils.equalsIgnoreCase(BAGGAGE_KEY_AGENT_TOKEN, entry.getKey())) {
                baggage.put(BAGGAGE_KEY_AGENT_TOKEN, entry.getValue());
            }
        }
        if (traceId != null && spanId != null) {
            return new TracerSpan(spanId, traceId, nextId(), baggage);
        }
        return new TracerSpan(NO_PARENT_ID, nextId(), nextId(), baggage);
    }

    public static TracerSpan extract(Http2Headers http2Headers) {
        String traceId = null;
        String spanId = null;
        Map<String, String> baggage = new HashMap<>();
        Iterator<Map.Entry<CharSequence, CharSequence>> entries = http2Headers.iterator();
        while (entries.hasNext()) {
            Map.Entry<CharSequence, CharSequence> entry = entries.next();
            String key = String.valueOf(entry.getKey());
            if (StringUtils.equalsIgnoreCase(BAGGAGE_ZIPKIN_TRACE_ID_KEY, key)) {
                traceId = String.valueOf(entry.getValue());
            } else if (StringUtils.equalsIgnoreCase(BAGGAGE_ZIPKIN_SPAN_ID_KEY, key)) {
                spanId = String.valueOf(entry.getValue());
            } else if (StringUtils.equalsIgnoreCase(BAGGAGE_KEY_ZIPKIN_SAMPLED, key)) {
                baggage.put(BAGGAGE_KEY_ZIPKIN_SAMPLED, String.valueOf(entry.getValue()));
            } else if (StringUtils.equalsIgnoreCase(BAGGAGE_KEY_AGENT_TOKEN, key)) {
                baggage.put(BAGGAGE_KEY_AGENT_TOKEN, String.valueOf(entry.getValue()));
            }
        }
        if (traceId != null && spanId != null) {
            return new TracerSpan(spanId, traceId, nextId(), baggage);
        }
        return new TracerSpan(NO_PARENT_ID, nextId(), nextId(), baggage);
    }

    public static void inject(TracerSpan span, HttpRequest httpRequest) {
        Iterator<Map.Entry<String, String>> entries = span.getBaggage().entrySet().iterator();
        while (entries.hasNext()) {
            Map.Entry<String, String> entry = entries.next();
            if (BAGGAGE_KEY_AGENT_TOKEN.equals(entry.getKey())
                    || BAGGAGE_KEY_ZIPKIN_SAMPLED.equals(entry.getKey())) {
                httpRequest.headers().set(entry.getKey(), entry.getValue());
            }
        }
        httpRequest.headers().set(BAGGAGE_ZIPKIN_SPAN_ID_KEY, span.getSpanId());
        httpRequest.headers().set(BAGGAGE_ZIPKIN_TRACE_ID_KEY, span.getTraceId());
    }

    public static void inject(TracerSpan span, Http2Headers http2Headers) {
        Iterator<Map.Entry<String, String>> entries = span.getBaggage().entrySet().iterator();
        while (entries.hasNext()) {
            Map.Entry<String, String> entry = entries.next();
            if (BAGGAGE_KEY_AGENT_TOKEN.equals(entry.getKey())) {
                http2Headers.set(BAGGAGE_KEY_AGENT_TOKEN_ASCII_KEY, Http2Util.newValue(entry.getValue()));
            } else if (BAGGAGE_KEY_ZIPKIN_SAMPLED.equals(entry.getKey())) {
                http2Headers.set(BAGGAGE_KEY_ZIPKIN_SAMPLED_ASCII_KEY, Http2Util.newValue(entry.getValue()));
            }
        }
        http2Headers.set(BAGGAGE_ZIPKIN_SPAN_ID_ASCII_KEY, Http2Util.newValue(span.getSpanId()));
        http2Headers.set(BAGGAGE_ZIPKIN_TRACE_ID_ASCII_KEY, Http2Util.newValue(span.getTraceId()));
    }

    public static String nextId() {
        return ZipkinIdGenerator.nextId();
    }


}
